#include "CRSolver.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
CRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::CRSolver(): IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>()
{
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Mult_ipl(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const {
    // Pipelined(b, x);
    // Chronopoulos_Gear(b, x);
    Standard(b, x);
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Standard(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const 
{
    assert(this->oper != nullptr);
    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    double prc_time = 0.0, op_time = 0.0, dot_time = 0.0;

    // 初始化辅助向量
    par_Vector<idx_t, ksp_t, dof> r(x);// residual
    par_Vector<idx_t, ksp_t, dof> u(x);// u := M^{-1}*r: preconditioned residual
    par_Vector<idx_t, ksp_t, dof> w(x);// w := A*u
    par_Vector<idx_t, ksp_t, dof> p(x);// search direction
    par_Vector<idx_t, ksp_t, dof> s(x);// s := A*p
    par_Vector<idx_t, ksp_t, dof> q(x);// q := M^{-1}*s
    r.set_val(0.0); u.set_val(0.0); w.set_val(0.0); p.set_val(0.0); s.set_val(0.0);
    ksp_t gamma, eta, alpha, beta;
    ksp_t tmp_loc[4], tmp_glb[4];// 用作局部点积的存储，发送缓冲区和接受缓冲区

    // compute initial residual
    if (this->zero_guess) {
        this->oper->Mult(x, r, false);
        vec_add(b, (ksp_t) -1.0, r, r);// r0 <- b-A*x0
    } else {
        vec_copy(b, r);
    }
    
    if (this->prec) {
        this->prec->Mult(r, u, true);// u0 <- M^{-1}*r0
    } else {
        vec_copy(r, u);
    }
    vec_copy(u, p);// 初始的搜索方向也就是预条件后的残差
    this->oper->Mult(p, s, false);
    vec_copy(s, w);// 初始的A*p 也就是A*M^{-1}*r=A*u
    if (this->prec) {
        this->prec->Mult(s, q, true);// q0 <- M^{-1}*s0
    } else {
        vec_copy(s, q);
    }

    tmp_loc[0] = tmp_loc[1] = tmp_loc[2] = tmp_loc[3] = 0.0;
    tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);// γ的局部点积
    tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(s, q);// η的局部点积
    tmp_loc[2] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(b, b);
    tmp_loc[3] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
    MPI_Allreduce(tmp_loc, tmp_glb, 4, MPI_DOUBLE, MPI_SUM, comm);
    gamma = tmp_glb[0];
    eta   = tmp_glb[1];
    ksp_t norm_b = sqrt(tmp_glb[2]);
    ksp_t norm_r = sqrt(tmp_glb[3]);
    alpha = gamma / eta;

    int & iter = this->final_iter = 0;
    if (my_pid == 0) printf("iter %4d   ||b|| = %.16e ||r||/||b|| = %.16e\n", iter, norm_b, norm_r / norm_b);

    iter++;// 初始化也算一步
    for ( ; iter < this->max_iter; iter++) {
        vec_add(x,  alpha, p, x);// x(i) <- x(i-1) + α(i-1)*p(i-1) 
        vec_add(r, -alpha, s, r);// r(i) <- r(i-1) - α(i-1)*s(i-1)
        vec_add(u, -alpha, q, u);// u(i) <- u(i-1) - α(i-1)*q(i-1)

        this->oper->Mult(u, w, false);// w(i) = A*u(i)

        double old_gamma = gamma;
        tmp_loc[0] = tmp_loc[1] = 0.0;
        tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);
        tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
        MPI_Allreduce(tmp_loc, tmp_glb, 2, MPI_DOUBLE, MPI_SUM, comm);
        gamma  = tmp_glb[0];
        norm_r = sqrt(tmp_glb[1]);

        double _rel = norm_r / norm_b;
        if (my_pid == 0) printf("iter %4d   alpha %.16e   ||r||/||b|| = %.16e\n", iter, alpha, _rel);
        if (_rel < this->rel_tol || norm_r < this->abs_tol) {// 判敛
            this->converged = 1;
            break;
        }

        beta = gamma / old_gamma;

        vec_add(u, beta, p, p);// p(i) <- u(i) + β(i) * p(i−1)
        vec_add(w, beta, s, s);// s(i) <- w(i) + β(i) * s(i-1)

        if (this->prec) {
            this->prec->Mult(s, q, true);// u(i) <- M^{-1}*r(i)
        } else {
            vec_copy(s, q);
        }

        eta = this->Dot(s, q);
        alpha = gamma / eta;
    }
    if (my_pid == 0) printf(" Std CR total iter: %d prec %.6f op %.6f dot %.6f s\n", iter, prc_time, op_time, dot_time);
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Chronopoulos_Gear(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const
{
    assert(this->oper != nullptr);
    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    double prc_time = 0.0, op_time = 0.0, dot_time = 0.0;

    // 初始化辅助向量
    par_Vector<idx_t, ksp_t, dof> r(x);// residual
    par_Vector<idx_t, ksp_t, dof> u(x);// u := M^{-1}*r: preconditioned residual
    par_Vector<idx_t, ksp_t, dof> w(x);// w := A*u
    par_Vector<idx_t, ksp_t, dof> v(x);// v := M^{-1}*w = M^{-1}*A*u
    par_Vector<idx_t, ksp_t, dof> p(x);// search direction
    par_Vector<idx_t, ksp_t, dof> s(x);// s := A*p
    par_Vector<idx_t, ksp_t, dof> q(x);// q := M^{-1}*s = M^{-1}*A*p
    r.set_val(0.0);
    u.set_val(0.0); w.set_val(0.0); v.set_val(0.0);
    p.set_val(0.0); s.set_val(0.0); q.set_val(0.0);

    ksp_t gamma, eta, alpha, beta;
    ksp_t tmp_loc[4], tmp_glb[4];// 用作局部点积的存储，发送缓冲区和接受缓冲区

    // compute initial residual
    if (this->zero_guess) {
        this->oper->Mult(x, r, false);
        vec_add(b, (ksp_t) -1.0, r, r);// r0 <- b-A*x0
    } else {
        vec_copy(b, r);
    }

    if (this->prec) {
        this->prec->Mult(r, u, true);// u0 <- M^{-1}*r0
    } else {
        vec_copy(r, u);
    }

    this->oper->Mult(u, w, false);

    if (this->prec) {
        this->prec->Mult(w, v, true);
    } else {
        vec_copy(w, v);
    }

    vec_copy(u, p);// 初始的搜索方向也就是预条件后的残差
    vec_copy(w, s);
    vec_copy(v, q);

    tmp_loc[0] = tmp_loc[1] = tmp_loc[2] = tmp_loc[3] = 0.0;
    tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);// γ>的局部点积
    tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(s, q);// η>的局部点积
    tmp_loc[2] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(b, b);
    tmp_loc[3] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
    MPI_Allreduce(tmp_loc, tmp_glb, 4, MPI_DOUBLE, MPI_SUM, comm);
    gamma = tmp_glb[0];
    eta   = tmp_glb[1];
    double norm_b = sqrt(tmp_glb[2]);
    double norm_r = sqrt(tmp_glb[3]);
    alpha = gamma / eta;

    int & iter = this->final_iter = 0;
    if (my_pid == 0) printf("iter %4d   ||b|| = %.16e ||r||/||b|| = %.16e\n", iter, norm_b, norm_r / norm_b);

    iter++;// 初始化也算一步
    for ( ; iter < this->max_iter; iter++) {
        vec_add(x,  alpha, p, x);// x(i) <- x(i-1) + α(i-1)*p(i-1) 
        vec_add(r, -alpha, s, r);// r(i) <- r(i-1) - α(i-1)*s(i-1)
        vec_add(u, -alpha, q, u);// u(i) <- u(i-1) - α(i-1)*q(i-1)

        this->oper->Mult(u, w, false);// w(i) = A*u(i)

         if (this->prec) {
            this->prec->Mult(w, v, true);// v(i) <- M^{-1}*w(i)
        } else {
            vec_copy(w, v);
        }

        double old_gamma = gamma;
        // TODO: kernel fusion
        tmp_loc[0] = tmp_loc[1] = tmp_loc[2] = tmp_loc[3] = 0.0;
        tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);
        tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(w, v);// 第一部分
        tmp_loc[2] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(s, v)
                    + vec_dot_intra_proc<idx_t, ksp_t, double, dof>(w, q);// 第二部分
        tmp_loc[3] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
        MPI_Allreduce(tmp_loc, tmp_glb, 4, MPI_DOUBLE, MPI_SUM, comm);
        gamma  = tmp_glb[0];
        // 要先更新β 才能计算η
        norm_r = sqrt(tmp_glb[3]);

        double _rel = norm_r / norm_b;
        if (my_pid == 0) printf("iter %4d   alpha %.16e   ||r||/||b|| = %.16e\n", iter, alpha, _rel);
        if (_rel < this->rel_tol || norm_r < this->abs_tol) {// 判敛
            this->converged = 1;
            break;
        }

        beta = gamma / old_gamma;

        vec_add(u, beta, p, p);// p(i) <- u(i) + β(i) * p(i−1)
        vec_add(w, beta, s, s);// s(i) <- w(i) + β(i) * s(i-1)
        vec_add(v, beta, q, q);// q(i) <- v(i) + β(i) * q(i-1)

        eta = tmp_glb[1] + beta * tmp_glb[2] + beta*beta * eta;
        alpha = gamma / eta;
    }
    if (my_pid == 0) printf(" CG-CR total iter: %d prec %.6f op %.6f dot %.6f s\n", iter, prc_time, op_time, dot_time);
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Pipelined(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const
{
    assert(this->oper != nullptr);
    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    double prc_time = 0.0, op_time = 0.0, dot_time = 0.0;

    // 初始化辅助向量
    par_Vector<idx_t, ksp_t, dof> r(x);// residual
    par_Vector<idx_t, ksp_t, dof> u(x);// u := M^{-1}*r: preconditioned residual
    par_Vector<idx_t, ksp_t, dof> w(x);// w := A*u
    par_Vector<idx_t, ksp_t, dof> v(x);// v := M^{-1}*w = M^{-1}*A*u
    par_Vector<idx_t, ksp_t, dof> p(x);// search direction
    par_Vector<idx_t, ksp_t, dof> s(x);// s := A*p
    par_Vector<idx_t, ksp_t, dof> q(x);// q := M^{-1}*s = M^{-1}*A*p
    par_Vector<idx_t, ksp_t, dof> f(x);// f := A*q
    par_Vector<idx_t, ksp_t, dof> h(x);// h := M^{-1}*f
    par_Vector<idx_t, ksp_t, dof> tmp_vec0(x), tmp_vec1(x);
    r.set_val(0.0);
    u.set_val(0.0); w.set_val(0.0); v.set_val(0.0);
    p.set_val(0.0); s.set_val(0.0); q.set_val(0.0);
    f.set_val(0.0); h.set_val(0.0); tmp_vec0.set_val(0.0); tmp_vec1.set_val(0.0);

    ksp_t gamma, eta, alpha, beta;
    ksp_t tmp_loc[4], tmp_glb[4];// 用作局部点积的存储，发送缓冲区和接受缓冲区

    // =========================== 初始化阶段 ======================================
    // compute initial residual
    if (this->zero_guess) {
        this->oper->Mult(x, r, false);
        vec_add(b, (ksp_t) -1.0, r, r);// r0 <- b-A*x0
    } else {
        vec_copy(b, r);
    }

    if (this->prec) {
        this->prec->Mult(r, u, true);// u0 <- M^{-1}*r0
    } else {
        vec_copy(r, u);
    }

    this->oper->Mult(u, w, false);
    if (this->prec) {
        this->prec->Mult(w, v, true);
    } else {
        vec_copy(w, v);
    }

    MPI_Request dot_req = MPI_REQUEST_NULL;
    tmp_loc[0] = tmp_loc[1] = tmp_loc[2] = tmp_loc[3] = 0.0;
    tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);// γ>的局部点积
    tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(w, v);// η>的局部点积
    tmp_loc[2] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(b, b);
    tmp_loc[3] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
    MPI_Iallreduce(tmp_loc, tmp_glb, 4, MPI_DOUBLE, MPI_SUM, comm, &dot_req);

    vec_copy(u, p);
    vec_copy(w, s);
    vec_copy(v, q);

    this->oper->Mult(q, f, false);
    if (this->prec) {
        this->prec->Mult(f, h, true);
    } else {
        vec_copy(f, h);
    }

    MPI_Wait(&dot_req, MPI_STATUS_IGNORE);
    gamma = tmp_glb[0];
    eta   = tmp_glb[1];
    ksp_t norm_b = sqrt(tmp_glb[2]);
    ksp_t norm_r = sqrt(tmp_glb[3]);
    alpha = gamma / eta;

    int & iter = this->final_iter = 0;
    if (my_pid == 0) printf("iter %4d   ||b|| = %.16e ||r||/||b|| = %.16e\n", iter, norm_b, norm_r / norm_b);

    iter++;// 初始化也算一步
    // ================================= 循环体 ===============================
    for ( ; iter < this->max_iter; iter++) {
        vec_add(x,  alpha, p, x);// x(i) <- x(i-1) + α(i-1)*p(i-1) 
        vec_add(r, -alpha, s, r);// r(i) <- r(i-1) - α(i-1)*s(i-1)
        vec_add(u, -alpha, q, u);// u(i) <- u(i-1) - α(i-1)*q(i-1)
        vec_add(w, -alpha, f, w);// w(i) <- w(i-1) - α(i-1)*f(i-1)
        vec_add(v, -alpha, h, v);// v(i) <- v(i-1) - α(i-1)*h(i-1)

        tmp_loc[0] = tmp_loc[1] = tmp_loc[2] = tmp_loc[3] = 0.0;
        tmp_loc[0] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, w);
        tmp_loc[1] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(w, v);// 第一部分
        tmp_loc[2] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(s, v)
                    + vec_dot_intra_proc<idx_t, ksp_t, double, dof>(w, q);// 第二部分
        tmp_loc[3] += vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
        
        MPI_Iallreduce(tmp_loc, tmp_glb, 4, MPI_DOUBLE, MPI_SUM, comm, &dot_req);

        this->oper->Mult(v, tmp_vec0, false);// t0 = A*v(i)
         if (this->prec) {
            this->prec->Mult(tmp_vec0, tmp_vec1, true);// t1 <- M^{-1}*t0
        } else {
            vec_copy(tmp_vec0, tmp_vec1);
        }

        MPI_Wait(&dot_req, MPI_STATUS_IGNORE);
        double old_gamma = gamma;
        gamma  = tmp_glb[0];
        // 要先更新β 才能计算η
        norm_r = sqrt(tmp_glb[3]);

        double _rel = norm_r / norm_b;
        if (my_pid == 0) printf("iter %4d   alpha %.16e   ||r||/||b|| = %.16e\n", iter, alpha, _rel);
        if (_rel < this->rel_tol || norm_r < this->abs_tol) {// 判敛
            this->converged = 1;
            break;
        }

        beta = gamma / old_gamma;

        vec_add(u, beta, p, p);// p(i) <- u(i) + β(i) * p(i−1)
        vec_add(w, beta, s, s);// s(i) <- w(i) + β(i) * s(i-1)
        vec_add(v, beta, q, q);// q(i) <- v(i) + β(i) * q(i-1)
        vec_add(tmp_vec0, beta, f, f);// f(i) <- t0 + β(i) * f(i-1)
        vec_add(tmp_vec1, beta, h, h);// h(i) <- t1 + β(i) * h(i-1)

        eta = tmp_glb[1] + beta * tmp_glb[2] + beta*beta * eta;
        alpha = gamma / eta;
    }
    if (my_pid == 0) printf(" Pipe-CR total iter: %d prec %.6f op %.6f dot %.6f s\n", iter, prc_time, op_time, dot_time);
}

template class CRSolver<int, double, double, double, 1>;
template class CRSolver<int, double, double, double, 2>;
template class CRSolver<int, double, double, double, 3>;
template class CRSolver<int, double, double, double, 4>;

// template class CRSolver<int, float , double, double>;
// template class CRSolver<int, float , float , double>;
// template class CRSolver<int, __fp16, float , double>;

// template class CRSolver<int, float , float , float >;
// template class CRSolver<int, __fp16, float , float >;
